Contents

Example 2: Understanding the inputs and the output trace

Example 2: Understanding the inputs and the output trace

In this notebook, we will discuss more in detail the inputs and the output trace of the VBMC algorithm, including a discussion on how to set parameter bounds and starting point. We also describe the items in the output trace, and display the evolution of the variational posterior in each iteration.

This notebook is part 1 of a series of notebooks in which we present various example usages for VBMC with the pyvbmc package.

import numpy as np
import scipy.stats as scs
from pyvbmc.vbmc import VBMC

1. Model definition: log likelihood, log prior, and log joint

We use the same toy target function as in Example 1, a broad Rosenbrock’s banana function in \(D = 2\).

D = 2  # still in 2-D

# log likelihood (Rosenbrock)
def llfun(x):
    x = np.atleast_2d(x)
    return -np.sum(
        (x[0, :-1] ** 2.0 - x[0, 1:]) ** 2.0 + (x[0, :-1] - 1) ** 2.0 / 100
    )

To make the problem more interesting, we will assume parameters to be constrained to be positive, that is \(x_d > 0\) for \(1 \le d \le D\).

Since parameters are strictly positive, we impose an independent exponential prior on each variable (again, you could choose here whatever you want). We then define the log joint (our target) as log likelihood plus log prior.

prior_tau = 3 * np.ones((1, D))  # Length scale of the exponential prior
lpriorfun = lambda x: np.sum(scs.expon.logpdf(x, scale=prior_tau))

target = lambda x: llfun(x) + lpriorfun(x)  # Log joint

2. Parameter setup: bounds and starting point

Choosing parameter bounds

Bound settings require some discussion:

  1. the specified (hard) bounds are not included in the domain, so in this case, since we want to be the parameters to be positive, we can set the lower bounds LB = 0, knowing that the parameter will always be greater than zero.

LB = np.zeros((1, D))  # Lower bounds

  1. Currently, VBMC does not support half-bounds, so we need to specify a finite upper bound (cannot be inf). We pick something very large according to our prior. However, do not go crazy here by picking something impossibly large, otherwise VBMC will likely fail.

UB = 10 * prior_tau  # Upper bounds

  1. Plausible bounds need to be meaningfully different from the hard bounds, so here we cannot set the plausible lower bounds PLB = 0. We follow the same strategy as the first example, by picking the top ~68% credible interval of the prior.

PLB = scs.expon.ppf(0.159, scale=prior_tau)
PUB = scs.expon.ppf(0.841, scale=prior_tau)

Choosing the starting point

Good practice is to initialize VBMC in a region of high posterior density.

  • With no additional information, uninformed approaches to choose the starting point \(\mathbf{x}_0\) include picking the mode of the prior, or, if you need multiple points, a random draw from the prior or a random point uniformly chosen inside the plausible box. None of these approaches is particularly recommended, though, especially when facing more difficult problems.

  • It is generally better to adopt an informed starting point \(\mathbf{x}_0\). A good choice is typically a point in the approximate vicinity of the mode of the posterior (no need to start exactly at the mode). Clearly, it is hard to know in advance the location of the mode, but a common approach is to run a few optimization steps on the target (starting from one of the uninformed options above).

For this example, we cheat and start in the proximity of the mode, exploiting our knowledge of the toy likelihood, for which we know the location of the maximum.

x0 = np.ones((1, D))  # Optimum of the Rosenbrock function

3. Initialize and run inference

Finally, we initialize and run VBMC.

Output trace

In each iteration, the VBMC output trace prints the following information:

  • Iteration: iteration number (starting from 0, which follows the initial design).

  • f-count: total number of the target function evaluations so far.

  • Mean[ELBO]: Estimated mean of the elbo.

  • Std[ELBO]: Standard deviation of the estimated elbo.

  • sKL-iter[q]: Gaussianized symmetrized Kullback-Leibler divergence between the variational posterior in the previous and current iteration. In brief, a measure of how much the variational posterior \(q\) changes between iterations.

  • K[q]: Number of mixture components of the variational posterior.

  • Convergence: A number that summarizes the stability of the current solution, by weighing various factors such as the change in the elbo and in the variational posterior across iterations, and the uncertainty in the elbo. Values (much) less than 1 correspond to a stable solution and thus suggest convergence of the algorithm.

  • Action: Major actions taken by the algorithm in the current iteration. Common actions/events are:

    • start warm-up and end warm-up: start and end of the initial warm-up stage;

    • trim data: some points are removed to increase stability;

    • stable: solution stable for multiple iteration;

    • finalize: refine the variational approximation with additional mixture components.

Iteration plotting

In this example, we also switch on the plotting of the current variational posterior at each iteration with the user option plot. Regarding the iteration plots:

  • Each iteration plot shows the multidimensional variational posterior in the current iteration via one- and two-dimensional marginal plots for each variable and pairs of variables, as you would get from vp.plot() (see also Example 1).

  • In each plot, blue circles represent points in the training set, and orange dots are points added to the training set in this iteration.

  • The red crosses are the centers of the variational mixture components.

  • Since VBMC with plotting on can occupy a lot of space on the notebook, remember that in Jupyter notebook you can toggle the scrolling on and off via the “Cell > Current Outputs > Toggle Scrolling” and “Cell > All Output > Toggle Scrolling” menus.

# initialize VBMC
user_options = {"plot": True}
vbmc = VBMC(target, x0, LB, UB, PLB, PUB, user_options)

# run VBMC
vp, elbo, elbo_sd, _, _ = vbmc.optimize()

Beginning variational optimization assuming EXACT observations of the log-joint.
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     0         10          -3.20         1.46    191163.68        2        inf     start warm-up
../_images/pyvbmc_example_2_16_1.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     1         15          -2.93         0.15         0.17        2        inf
../_images/pyvbmc_example_2_16_3.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     2         20          -2.61         0.17         0.09        2       3.74
../_images/pyvbmc_example_2_16_5.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     3         25          -2.53         0.03         0.06        2        1.8
../_images/pyvbmc_example_2_16_7.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     4         30          -2.09         1.99         0.53        2       20.6     end warm-up
../_images/pyvbmc_example_2_16_9.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     5         35          -2.40         0.05         0.14        2       4.38
../_images/pyvbmc_example_2_16_11.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     6         40          -2.38         0.05         0.04        2       1.18
../_images/pyvbmc_example_2_16_13.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     7         45          -2.19         0.04         0.07        3       2.51
../_images/pyvbmc_example_2_16_15.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     8         50          -2.17         0.01         0.02        4      0.473
../_images/pyvbmc_example_2_16_17.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
     9         55          -2.14         0.01         0.02        7       0.57
../_images/pyvbmc_example_2_16_19.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    10         60          -2.09         0.01         0.01       10      0.377
../_images/pyvbmc_example_2_16_21.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    11         65          -2.05         0.01         0.01       13      0.328
../_images/pyvbmc_example_2_16_23.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    12         70          -2.01         0.03         0.01       14      0.375
../_images/pyvbmc_example_2_16_25.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    13         75          -2.02         0.00         0.00       14      0.183
../_images/pyvbmc_example_2_16_27.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    14         80          -1.99         0.00         0.00       15      0.116
../_images/pyvbmc_example_2_16_29.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    15         85          -2.01         0.00         0.00       16     0.0625
../_images/pyvbmc_example_2_16_31.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    16         90          -1.99         0.00         0.00       15     0.0762
../_images/pyvbmc_example_2_16_33.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
    17         95          -1.98         0.00         0.00       14      0.046     stable
../_images/pyvbmc_example_2_16_35.png 
 Iteration  f-count    Mean[ELBO]    Std[ELBO]    sKL-iter[q]   K[q]  Convergence  Action
   inf         95          -1.93         0.01         0.00       50      0.046     finalize
../_images/pyvbmc_example_2_16_37.png 
Inference terminated: variational solution stable for options.tolstablecountfcn evaluations.
Estimated ELBO: -1.926 +/-0.006.

lml_true = -1.836  # ground truth, which we know for this toy scenario

print("The true log model evidence is:", lml_true)
print("The obtained ELBO is:", format(elbo, ".3f"))
print("The obtained ELBO_SD is:", format(elbo_sd, ".3f"))

The true log model evidence is: -1.836
The obtained ELBO is: -1.926
The obtained ELBO_SD is: 0.006

4. Examine and visualize results

We will examine now the obtained variational posterior with an interactive plot.

Note: The interactive plot will not be visible in a notebook preview on GitHub. See below for a static image.

# Generate samples from the variational posterior
n_samples = int(3e5)
Xs, _ = vp.sample(n_samples)

# We compute the pdf of the approximate posterior on a 2-D grid
plot_lb = np.zeros(2)
plot_ub = np.quantile(Xs, 0.999, axis=0)
x1 = np.linspace(plot_lb[0], plot_ub[0], 400)
x2 = np.linspace(plot_lb[1], plot_ub[1], 400)

xa, xb = np.meshgrid(x1, x2)  # Build the grid
xx = np.vstack(
    (xa.ravel(), xb.ravel())
).T  # Convert grids to a vertical array of 2-D points
yy = vp.pdf(xx)  # Compute PDF values on specified points

# Plot approximate posterior pdf (this interactive plot does not work in higher D)
import plotly.graph_objects as go

fig = go.Figure(
    data=[
        go.Surface(z=yy.reshape(x1.size, x2.size), x=x1, y=x2, showscale=False)
    ]
)
fig.update_layout(
    autosize=False,
    width=800,
    height=500,
    scene=dict(
        xaxis_title="x_0",
        yaxis_title="x_1",
        zaxis_title="Approximate posterior pdf",
    ),
)

# Compute and plot approximate posterior mean
post_mean = np.mean(Xs, axis=0)
fig.add_trace(
    go.Scatter3d(
        x=[post_mean[0]],
        y=[post_mean[1]],
        z=vp.pdf(post_mean),
        name="Approximate posterior mean",
        marker_symbol="diamond",
    )
)

# Find and plot approximate posterior mode
# post_mode = vp.mode()
# fig.add_trace(go.Scatter3d(x=[post_mode[0]], y=[post_mode[1]], z=vp.pdf(post_mode), name="Approximate posterior mode", marker_symbol="x"))

# Display posterior also as a non-interactive static image
from IPython.display import Image

Image(fig.to_image(format="png"))
../_images/pyvbmc_example_2_21_0.png

5. Conclusions

In this notebook, we have discussed more in detail the input parameter settings (bounds and starting point) and the output trace of VBMC.

An important step still not considered in this example is a thorough validation of the results and diagnostics, which will be discussed in a later notebook.

Example 2: full code

The following cell includes in a single place all the code used in Example 2, without the extra fluff.

import numpy as np
import scipy.stats as scs
from pyvbmc.vbmc import VBMC

#####################################################################
# 1. Model definition: log likelihood, log prior, and log joint

D = 2  # still in 2-D

# log likelihood (Rosenbrock)
def llfun(x):
    x = np.atleast_2d(x)
    return -np.sum(
        (x[0, :-1] ** 2.0 - x[0, 1:]) ** 2.0 + (x[0, :-1] - 1) ** 2.0 / 100
    )


prior_tau = 3 * np.ones((1, D))  # Length scale of the exponential prior
lpriorfun = lambda x: np.sum(scs.expon.logpdf(x, scale=prior_tau))

target = lambda x: llfun(x) + lpriorfun(x)  # Log joint

#####################################################################
# 2. Parameter setup: bounds and starting point

LB = np.zeros((1, D))  # Lower bounds
UB = 10 * prior_tau  # Upper bounds
PLB = scs.expon.ppf(0.159, scale=prior_tau)
PUB = scs.expon.ppf(0.841, scale=prior_tau)

x0 = np.ones((1, D))  # Optimum of the Rosenbrock function

#####################################################################
# 3. Initialize and run inference

# initialize VBMC
user_options = {"plot": True}
vbmc = VBMC(target, x0, LB, UB, PLB, PUB, user_options)

# run VBMC
vp, elbo, elbo_sd, _, _ = vbmc.optimize()

lml_true = -1.836  # ground truth, which we know for this toy scenario
print("The true log model evidence is:", lml_true)
print("The obtained ELBO is:", format(elbo, ".3f"))
print("The obtained ELBO_SD is:", format(elbo_sd, ".3f"))

#####################################################################
# 4. Analyze and visualize results

# Generate samples from the variational posterior
n_samples = int(3e5)
Xs, _ = vp.sample(n_samples)

# We compute the pdf of the approximate posterior on a 2-D grid
plot_lb = np.zeros(2)
plot_ub = np.quantile(Xs, 0.999, axis=0)
x1 = np.linspace(plot_lb[0], plot_ub[0], 400)
x2 = np.linspace(plot_lb[1], plot_ub[1], 400)

xa, xb = np.meshgrid(x1, x2)  # Build the grid
xx = np.vstack(
    (xa.ravel(), xb.ravel())
).T  # Convert grids to a vertical array of 2-D points
yy = vp.pdf(xx)  # Compute PDF values on specified points

# Plot approximate posterior pdf (this interactive plot does not work in higher D)
import plotly.graph_objects as go

fig = go.Figure(
    data=[
        go.Surface(z=yy.reshape(x1.size, x2.size), x=x1, y=x2, showscale=False)
    ]
)
fig.update_layout(
    autosize=False,
    width=800,
    height=500,
    scene=dict(
        xaxis_title="x_0",
        yaxis_title="x_1",
        zaxis_title="Approximate posterior pdf",
    ),
)

# Compute and plot approximate posterior mean
post_mean = np.mean(Xs, axis=0)
fig.add_trace(
    go.Scatter3d(
        x=[post_mean[0]],
        y=[post_mean[1]],
        z=vp.pdf(post_mean),
        name="Approximate posterior mean",
        marker_symbol="diamond",
    )
)

# Find and plot approximate posterior mode
# post_mode = vp.mode()
# fig.add_trace(go.Scatter3d(x=[post_mode[0]], y=[post_mode[1]], z=vp.pdf(post_mode), name="Approximate posterior mode", marker_symbol="x"))

# Display posterior also as a non-interactive static image
from IPython.display import Image

Image(fig.to_image(format="png"))

Acknowledgments

Work on the pyvbmc package was funded by the Finnish Center for Artificial Intelligence FCAI.

previous

Example 1: Basic usage

next

About us